package process

import (
	"context"
	"encoding/json"
	llmMaker "quipus/internal/llm/factory"
	"quipus/internal/rag/rag"
	ragStore "quipus/internal/rag/storage"
	"quipus/internal/utils"
	"strconv"

	"github.com/spf13/viper"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

type RagConfig struct {
	VectorDB   ragStore.VectorDBConfig `yaml:"vectordb"`
	GraphDB    ragStore.GraphDBConfig  `yaml:"graphdb"`
	MaxToken   int                     `yaml:"max_token"`
	ChunkToken int                     `yaml:"chunk_token"`
	ScoreBias  float32                 `yaml:"score_bias"`
}

var (
	ragConfig = RagConfig{}
)

func init() {
	utils.ReadYamlConfig("etc/rag/rag.yaml", &ragConfig)
}

func (p *Process) NewRagBuilder(config *rag.RagClientConfig) (*rag.RagBuilder, error) {

	ctx, cancel := context.WithCancel(context.Background())

	ragBuilder, e := rag.NewRagBuilder(ctx, config)
	if e != nil {
		return nil, status.Errorf(codes.Internal, "failed to create rag builder: %v", e)
	}
	ragBuilder.Cancel = cancel
	ragBuilder.Ctx = ctx
	return ragBuilder, nil
}

func (p *Process) SetRagAssistant(llmId int32, emId int32, ragConfig *rag.RagClientConfig) (*rag.RagBuilder, error) {
	ctx := context.Background()

	lp, lm, err := p.GetModelProvider(ctx, llmId)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get assistant: %v", err)
	}

	llmConfig := &llmMaker.ChatClientConfig{
		Url:   lp.AccessURL,
		Model: lm.Model,
		API:   lp.APIName,
	}

	if lm.OptionInfo != "" {
		option := make(map[string]string)
		err := json.Unmarshal([]byte(lm.OptionInfo), &option)
		if err != nil {
			return nil, status.Errorf(codes.Internal, "failed to unmarshal option info: %v", err)
		}
		llmConfig.MaxTokens, err = strconv.Atoi(option["MaxTokens"])
		if err != nil {
			return nil, status.Errorf(codes.Internal, "failed to convert dim: %v", err)
		}
	}

	if lp.AuthInfo != "" {
		key := []byte(viper.GetString("encryption-key"))
		authInfo, err := utils.Decrypt(lp.AuthInfo, key)
		if err != nil {
			return nil, status.Errorf(codes.Internal, "failed to decrypt auth info: %v", err)
		}
		llmConfig.AuthInfo = authInfo
	}

	ragConfig.ChatConfig = llmConfig
	ep, em, err := p.GetModelProvider(ctx, emId)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get assistant: %v", err)
	}

	embedConfig := &llmMaker.EmbeddingClientConfig{
		Url:   ep.AccessURL,
		Model: em.Model,
		API:   ep.APIName,
	}

	if em.OptionInfo != "" {
		option := make(map[string]string)
		err := json.Unmarshal([]byte(em.OptionInfo), &option)
		if err != nil {
			return nil, status.Errorf(codes.Internal, "failed to unmarshal option info: %v", err)
		}
		embedConfig.Dim, err = strconv.Atoi(option["Dim"])
		if err != nil {
			return nil, status.Errorf(codes.Internal, "failed to convert dim: %v", err)
		}
	}

	if ep.AuthInfo != "" {
		key := []byte(viper.GetString("encryption-key"))
		authInfo, err := utils.Decrypt(ep.AuthInfo, key)
		if err != nil {
			return nil, status.Errorf(codes.Internal, "failed to decrypt auth info: %v", err)
		}
		embedConfig.AuthInfo = authInfo
	}

	ragConfig.EmbeddingConfig = embedConfig
	return p.NewRagBuilder(ragConfig)
}

func (p *Process) NewRagConfig() *rag.RagClientConfig {

	config := &rag.RagClientConfig{
		VectorConfig: &ragConfig.VectorDB,
		GraphConfig:  &ragConfig.GraphDB,
		MaxTokens:    ragConfig.MaxToken,
		ChunkTokens:  ragConfig.ChunkToken,
		ScoreBias:    ragConfig.ScoreBias,
	}

	return config
}
